/**
    Copyright (c) 2010 yakiimo02
    Distributed under the New BSD License.
    See included license.txt or http://www.yakiimo3d.com/NewBSDLicense.txt
**/

/**
    @file   
    @brief  Buddhabrot compute shader.
    @author yakiimo02
    @date    2010/01/23
*/

Texture2D randomInput                                : register( t0 );                // random number input texture
RWStructuredBuffer< uint4 >        iterationsBuffer    : register (u0);                // buddhabrot & nebulabrot iteration plots

cbuffer cb : register( b0 )
{
    // bounds of the complex number plane
    float    fMinRe;
    float    fMaxRe;
    float    fMinIm;
    float    fMaxIm;
    
    float    fMinReZoom;
    float    fMaxReZoom;
    float    fMinImZoom;
    float    fMaxImZoom;
    
    float    fReC;                // precalculated const
    float    fImC;                // precalculated const
    
    // input dim
    int        nInputWidth;
    int        nInputHeight;

    // output dim
    int        nOutputWidth;
    int        nOutputHeight;

    int        nMaxIteration; // largest max iteration among the 3 nebulabrot rgb maxIterations

    // nebulabrot max iterations
    int        maxIterationsR;
    int        maxIterationsG;
    int        maxIterationsB;
};

/**
    Iterate and determine whether a complex number is a member of the mandelbrot set.
    @param zRe    [in]        the real part of a complex number
    @param zIm    [in]        the imaginary part of a complex number
*/
int Iterate( float zRe, float zIm ) 
{
    float cRe = zRe;
    float cIm = zIm;

    for( int nIteration=0; nIteration<nMaxIteration; ++nIteration )
    {
        float zRe2 = zRe*zRe;
        float zIm2 = zIm*zIm;

        /// check whether absolute value of z has gotten bigger than 2
        if( zRe2 + zIm2 > 4.0 ) {
            return nIteration;
        }

        // calculate next z(n)
        zIm = 2.0*zRe*zIm + cIm;
        zRe = zRe2 - zIm2 + cRe;
    }

    // complex number did not escape to infinity 
    return nMaxIteration;
}

/**
    Iterate and plot the path of the escaping complex number.
    @param zRe                [in]        the real part of a complex number
    @param zIm                [in]        the imaginary part of a complex number
    @param nWidth            [in]        width of the output buffer
    @param nHeight            [in]        height of the output buffer
    @param nIterationEscape    [in]        the maximum iterations to plot    
*/
void IterateAndPlot( float zRe, float zIm, int nWidth, int nHeight, int nIterationEscape ) 
{
    float cRe = zRe;
    float cIm = zIm;
    
    for( int nIteration=0; nIteration<=nIterationEscape; ++nIteration )
    {
        //int nX = ( int )( nWidth * (zRe - fMinReZoom) / (fMaxReZoom-fMinReZoom) );
        //int nY = ( int )( nHeight * (zIm - fMinImZoom) / (fMaxImZoom-fMinImZoom) );
        int nX = (zRe - fMinReZoom) * fReC;
        int nY = (zIm - fMinImZoom) * fImC;
            
        if( nX >=0 && nY >=0 && nX < nWidth && nY < nHeight ) 
        {
            int nIndex = nX + nWidth*nY;
        #if 1
            if( nIteration < maxIterationsR ) {
                InterlockedAdd( iterationsBuffer[ nIndex ].x, 1 );
            }
            if( nIteration < maxIterationsG ) {
                InterlockedAdd( iterationsBuffer[ nIndex ].y, 1 );
            }
            if( nIteration < maxIterationsB ) {
                InterlockedAdd( iterationsBuffer[ nIndex ].z, 1 );
            }
        // Check performance for when the output buffer isn't synchronized. Results are wrong
        // but wanted to check speed. I got the odd result that using synchronizing InterlockedAdd 
        // is faster.
        #else
            if( nIteration < maxIterationsR ) {
                iterationsBuffer[ nIndex ].x += 1;
            }
            if( nIteration < maxIterationsG ) {
                iterationsBuffer[ nIndex ].y += 1;
            }
            if( nIteration < maxIterationsB ) {
                iterationsBuffer[ nIndex ].z += 1;
            }
        #endif
        }

        float zRe2 = zRe*zRe;
        float zIm2 = zIm*zIm;

        // calculate next z(n)
        zIm = 2.0*zRe*zIm + cIm;
        zRe = zRe2 - zIm2 + cRe;
    }
}

// (10,10,1) thread groups for 10,000 threads per frame
[numthreads(10, 10, 1)]
void BuddhabrotCS( uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex )
{
    float x = DTid.x;
    float y = DTid.y;

    int nInputIndex = DTid.x+DTid.y*nInputWidth;
    
    // sample a random number from the complex plane
    float2 random = randomInput.Load( uint3(DTid.x, DTid.y, 0) ).xy;
    float zRe = random.x;
    float zIm = random.y;
        
    // check whether sample complex number is in the mandelbrot set.
    int nIterationEscape = Iterate( zRe, zIm );

    if( nIterationEscape != nMaxIteration ) { 
        // If the complex number is not a member of the mandelbrot set, plot the complex number's path as
        // it escapes to infinity. 
        IterateAndPlot( zRe, zIm, nOutputWidth, nOutputHeight, nIterationEscape );        
    }    
}
